import numpy as np
import pandas as pd
import funs as func_


'''
Example code to calculate nonparametric conditional independence test
See the function FUNS for details on the implementation and usage
'''

NOBS = 1000
np.random.seed(5)
#%% Example of conditional independence on a fork
#=== Create synthetic data 
t, conditionining_set_size = func_.FORK(NOBS)

#=== Test for unconditional independence X _||_ Y 
#=== In the function below X = 0 (the first column), Y = 1 (the second columns), None = conditioning set
unconditional = func_.independence_test(t, 0, 1, None, random_seed = True)
unlik = ['likely dependend' if unconditional < 0.1 else 'likely independent'][0]
#=== Test for the conditional independence X _||_ Y | Z
conditional = func_.independence_test(t, 0, 1, [2+n for n in range(conditionining_set_size)], random_seed = True)
colik = ['likely dependend' if conditional < 0.1 else 'likely independent'][0]



print('Fork X <- Z -> Y, then X and Y are:\nUnconditional test: ', np.round(unconditional, 2), ' (', unlik, ')', '\nConditional test: ', np.round(conditional,2), ' (', colik, ')', sep='')


#%% Example of conditional independence on a chain
#=== Create synthetic data 
t, conditionining_set_size = func_.CHAIN(NOBS)

#=== Test for unconditional independence X _||_ Y 
#=== In the function below X = 0 (the first column), Y = 1 (the second columns), None = conditioning set
unconditional = func_.independence_test(t, 0, 1, None, random_seed = True)
unlik = ['likely dependend' if unconditional < 0.1 else 'likely independent'][0]
#=== Test for the conditional independence X _||_ Y | Z
conditional = func_.independence_test(t, 0, 1, [2+n for n in range(conditionining_set_size)], random_seed = True)
colik = ['likely dependend' if conditional < 0.1 else 'likely independent'][0]


print('Chain X -> Z -> Y, then X and Y are\nUnconditional test: ', np.round(unconditional, 2), ' (', unlik, ')', '\nConditional test: ', np.round(conditional,2), ' (', colik, ')', sep='')


#%% Example of conditional independence on a collider
#=== Create synthetic data 
t, conditionining_set_size = func_.COLLIDER(NOBS)


#=== Test for unconditional independence X _||_ Y 
#=== In the function below X = 0 (the first column), Y = 1 (the second columns), None = conditioning set
unconditional = func_.independence_test(t, 0, 1, None, random_seed = True)
unlik = ['likely dependend' if unconditional < 0.1 else 'likely independent'][0]
#=== Test for the conditional independence X _||_ Y | Z
conditional = func_.independence_test(t, 0, 1, [2+n for n in range(conditionining_set_size)], random_seed = True)
colik = ['likely dependend' if conditional < 0.1 else 'likely independent'][0]



print('Collider X -> Z <- Y, then X and Y are\nUnconditional test: ', np.round(unconditional, 2), ' (', unlik, ')', '\nConditional test: ', np.round(conditional,2), ' (', colik, ')', sep='')
